#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 改动要点：
# 1) joystick 轴做 EMA (low-pass) + deadzone
# 2) 最小位移/旋转阈值，超过才更新 target_pose
# 3) IK 结果与当前关节角比较：只有大于 joint_send_threshold_deg 才下发 SDK
# 4) 下发最小时间间隔 throttling
# 5) 从 SDK 读回 joint_states 做简单低通，发布平滑后的 joint_states

import rospy
import DrEmpower as dr
import math
import threading
import numpy as np
import tf.transformations as tft
from sensor_msgs.msg import Joy, JointState
from geometry_msgs.msg import PoseStamped
from moveit_commander import MoveGroupCommander, roscpp_initialize, roscpp_shutdown
from moveit_msgs.srv import GetPositionIK, GetPositionIKRequest, GetPositionFK, GetPositionFKRequest
from copy import deepcopy
import time

# ---------------- 配置区 ----------------
joint_names = ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6"]
joint_ids   = [1, 2, 3, 4, 5, 6]
# ----------------------------------------

class TeleopPoseReal:
    def __init__(self):
        roscpp_initialize([])
        rospy.init_node("real_remote_full_6d")

        # 参数（可在 launch/param 设置覆盖）
        self.group_name = rospy.get_param("~group_name", "robot_arm")
        self.step_xyz = rospy.get_param("~step_xyz", 0.002)
        self.step_rpy = rospy.get_param("~step_rpy", 0.01)
        self.loop_rate = rospy.get_param("~loop_rate", 50)
        self.joy_deadzone = rospy.get_param("~joy_deadzone", 0.12)  # 增大死区
        self.axis_ema_alpha = rospy.get_param("~axis_ema_alpha", 0.35)  # EMA 平滑系数 (0..1)
        self.translation_min_move = rospy.get_param("~translation_min_move", 0.0008)  # m 最小平移
        self.rotation_min_move = rospy.get_param("~rotation_min_move", 0.3 * math.pi/180.0)  # rad 最小旋转 (~0.3°)
        self.joint_send_threshold_deg = rospy.get_param("~joint_send_threshold_deg", 0.5)  # 关节角差阈值 (deg)
        self.min_send_interval = rospy.get_param("~min_send_interval", 0.08)  # s 下发最小间隔
        self.joint_states_smooth_alpha = rospy.get_param("~joint_states_smooth_alpha", 0.25)

        # MoveIt 初始化
        self.group = MoveGroupCommander(self.group_name)
        self.ee_link = self.group.get_end_effector_link()
        self.joint_names = joint_names

        # 初始化 joystick 状态（包含 EMA 状态）
        self.joy_axes = [0.0] * 8
        self.joy_axes_ema = [0.0] * 8

        # 初始化 joint_states（读取实物角度）
        rospy.loginfo("🔹 正在读取机械臂关节角度...")
        positions_rad = []
        for jid in joint_ids:
            try:
                ang_deg = self.safe_dr_get_angle(jid)
                positions_rad.append(math.radians(ang_deg))
            except Exception as e:
                rospy.logwarn(f"读取关节 {jid} 失败: {e}")
                positions_rad.append(0.0)
        rospy.loginfo(f"✅ 读取到关节角度 (rad): {positions_rad}")

        # 使用 FK 求初始末端位姿
        fk_req = GetPositionFKRequest()
        fk_req.header.frame_id = "base_link"
        fk_req.fk_link_names = [self.ee_link]
        js_init = JointState()
        js_init.name = joint_names
        js_init.position = positions_rad
        fk_req.robot_state.joint_state = js_init
        try:
            fk_res = rospy.ServiceProxy("/compute_fk", GetPositionFK)(fk_req)
            self.target_pose = PoseStamped()
            self.target_pose.header.frame_id = "base_link"
            self.target_pose.pose = fk_res.pose_stamped[0].pose
        except Exception as e:
            rospy.logwarn(f"FK初始化失败: {e}")
            current_pose = self.group.get_current_pose()
            self.target_pose = PoseStamped()
            self.target_pose.header = current_pose.header
            self.target_pose.pose = current_pose.pose

        self.print_pose(self.target_pose, "初始位姿（实物 FK）")

        # ROS 通信
        rospy.Subscriber("/joy", Joy, self.cb_joy)
        self.joint_state_pub = rospy.Publisher("/joint_states", JointState, queue_size=10)

        # MoveIt 服务
        rospy.wait_for_service("/compute_ik")
        rospy.wait_for_service("/compute_fk")
        self.ik_srv = rospy.ServiceProxy("/compute_ik", GetPositionIK)
        self.fk_srv = rospy.ServiceProxy("/compute_fk", GetPositionFK)

        # 互斥锁
        self.lock = threading.Lock()

        # 启动 joint_states 发布线程（平滑读回）
        self.joint_pub_thread = threading.Thread(target=self.publish_joint_states_loop, daemon=True)
        self.joint_pub_thread.start()

        # 发送/下发节流
        self._last_send_time = rospy.Time(0)

        rospy.loginfo("✅ 实物遥操作节点启动（含抖动抑制）")
        self.run()

    # ---------- 手柄回调（做 EMA 平滑 + 死区 + 按键动作）----------
    def cb_joy(self, msg: Joy):
        # ---- 1) 轴平滑 ----
        for i in range(min(len(self.joy_axes), len(msg.axes))):
            raw = msg.axes[i]
            ema_prev = self.joy_axes_ema[i]
            ema_new = (1.0 - self.axis_ema_alpha) * ema_prev + self.axis_ema_alpha * raw
            self.joy_axes_ema[i] = ema_new
            self.joy_axes[i] = ema_new if abs(ema_new) > self.joy_deadzone else 0.0

        # ---- 2) 按钮检测（新增）----
        # 初始化前一帧按钮状态缓存
        if not hasattr(self, "last_buttons"):
            self.last_buttons = [0] * len(msg.buttons)

        # 检测按钮边沿（按下瞬间触发）
        for i, val in enumerate(msg.buttons):
            if val == 1 and self.last_buttons[i] == 0:
                # 按钮按下沿触发
                if i == 0:
                    rospy.loginfo("🟢 Button[0] pressed → 机械臂复位中 ...")
                    threading.Thread(target=self.goto_joint_angles, args=([0, 0, 0, 0, 0, 0], "复位"), daemon=True).start()
                elif i == 1:
                    rospy.loginfo("🟢 Button[1] pressed → 机械臂前往观测点 ...")
                    threading.Thread(target=self.goto_joint_angles, args=([35.11, 24.06, 42.97, -38.97, -90, 0], "观测点"), daemon=True).start()

        # 更新按钮缓存
        self.last_buttons = msg.buttons[:]

    # ---------- 按钮动作：直接下发目标角度 ----------
    def goto_joint_angles(self, target_deg_list, name=""):
        with self.lock:
            try:
                dr.set_angles(id_list=joint_ids, angle_list=target_deg_list, speed=20, param=10, mode=1)
                rospy.loginfo(f"✅ 机械臂已移动至{name}位置: {target_deg_list}")
                # 同步更新 target_pose
                positions_rad = [math.radians(a) for a in target_deg_list]
                fk_req = GetPositionFKRequest()
                fk_req.header.frame_id = "base_link"
                fk_req.fk_link_names = [self.ee_link]
                js = JointState()
                js.name = self.joint_names
                js.position = positions_rad
                fk_req.robot_state.joint_state = js
                fk_res = self.fk_srv(fk_req)
                self.target_pose.pose = fk_res.pose_stamped[0].pose
            except Exception as e:
                rospy.logerr(f"⚠️ 无法移动至{name}位置: {e}")


    # ---------- 主循环 ----------
    def run(self):
        rate = rospy.Rate(self.loop_rate)
        while not rospy.is_shutdown():
            if any(abs(v) > 1e-6 for v in self.joy_axes):
                prev = np.array([
                    self.target_pose.pose.position.x,
                    self.target_pose.pose.position.y,
                    self.target_pose.pose.position.z
                ])
                prev_q = deepcopy(self.target_pose.pose.orientation)

                # 先尝试更新目标位姿，但只有在超过最小移动阈值时才视为“有效更新”
                updated = self.update_target_pose_with_threshold(prev, prev_q)
                if updated:
                    if not self.solve_and_publish():
                        self.rollback_pose(prev, prev_q)
            rate.sleep()

    # ---------- 位姿更新（只有大于阈值时改变 target_pose） ----------
    def update_target_pose_with_threshold(self, prev_pos, prev_q):
        # 计算目标增量但不直接应用：判断是否超过最小阈值
        dx, dy, dz = [self.step_xyz * self.joy_axes[i] for i in [0, 3, 1]]
        droll = self.step_rpy * self.joy_axes[7]
        dpitch = self.step_rpy * self.joy_axes[2]
        dyaw = self.step_rpy * self.joy_axes[6]

        # 小增量都被忽略（避免抖动），但以当前姿态为基准计算实际世界增量
        if abs(dx) < 1e-12 and abs(dy) < 1e-12 and abs(dz) < 1e-12 and abs(droll) < 1e-12 and abs(dpitch) < 1e-12 and abs(dyaw) < 1e-12:
            return False

        q = self.target_pose.pose.orientation
        rot = self.quat_to_rot(q)
        delta = rot.dot(np.array([dx, dy, dz]))
        trans_magnitude = np.linalg.norm(delta)
        rot_magnitude = max(abs(droll), abs(dpitch), abs(dyaw))

        # 若平移/旋转都小于阈值，则忽略
        if trans_magnitude < self.translation_min_move and rot_magnitude < self.rotation_min_move:
            return False

        # 应用变化
        p = self.target_pose.pose.position
        p.x += delta[0]
        p.y += delta[1]
        p.z += delta[2]

        dq = tft.quaternion_from_euler(droll, dpitch, dyaw)
        q_now = [q.x, q.y, q.z, q.w]
        q_new = tft.quaternion_multiply(q_now, dq)
        self.target_pose.pose.orientation.x = q_new[0]
        self.target_pose.pose.orientation.y = q_new[1]
        self.target_pose.pose.orientation.z = q_new[2]
        self.target_pose.pose.orientation.w = q_new[3]

        self.print_pose(self.target_pose, "目标位姿（更新）", throttle=0.5)
        return True

    # ---------- IK + 校验（并且节流下发） ----------
    def solve_and_publish(self):
        # 节流：不要每帧都请求 IK（但这里只是防止短时间内重复下发，实际上 IK 请求可以保留）
        now = rospy.Time.now()
        # 构造请求
        req = GetPositionIKRequest()
        req.ik_request.group_name = self.group_name
        req.ik_request.pose_stamped = self.target_pose
        req.ik_request.ik_link_name = self.ee_link
        req.ik_request.timeout = rospy.Duration(0.05)
        req.ik_request.robot_state.joint_state.name = self.joint_names
        req.ik_request.robot_state.joint_state.position = self.get_current_joint_values()
        try:
            res = self.ik_srv(req)
            if res.error_code.val != 1:
                rospy.logwarn_throttle(1.0, f"❌ IK失败 code={res.error_code.val}")
                return False
        except Exception as e:
            rospy.logerr(f"IK服务错误: {e}")
            return False

        # 检查与当前读回关节差异：若小于阈值就不下发以避免抖动
        sol_positions = list(res.solution.joint_state.position[:6])
        current_positions = self.get_current_joint_values()
        # 转为度
        diffs_deg = [abs(math.degrees(sol_positions[i] - current_positions[i])) for i in range(6)]
        max_diff = max(diffs_deg)
        if max_diff < self.joint_send_threshold_deg:
            rospy.logdebug(f"下发被跳过（关节差小）：max_diff={max_diff:.3f} deg")
            return True  # 认为成功但不发送

        # 节流下发：确保最小发送间隔
        if (now - self._last_send_time).to_sec() < self.min_send_interval:
            rospy.logdebug("下发被节流（发送频率过高）")
            return True

        # 实际发送：转为度并调用 SDK
        self.publish_to_sdk(res.solution.joint_state)
        self._last_send_time = rospy.Time.now()
        rospy.loginfo_throttle(0.5, f"✅ 已发送 IK 结果至实物机械臂（max_diff={max_diff:.2f}°）")
        return True

    # ---------- 发送到实物（保持不变，但捕获异常） ----------
    def publish_to_sdk(self, joint_state):
        joint_deg = [np.degrees(a) for a in joint_state.position[:6]]
        with self.lock:
            try:
                # 注意：speed/param/mode 根据你的底层 SDK 文档调整
                dr.set_angles(id_list=joint_ids, angle_list=joint_deg, speed=20, param=10, mode=1)
            except Exception as e:
                rospy.logerr(f"⚠️ SDK发送失败: {e}")

    # ---------- joint_states 循环发布（读取并平滑） ----------
    def publish_joint_states_loop(self):
        rate = rospy.Rate(50)
        # 初始化平滑缓存
        smooth_positions = [0.0] * 6
        initialized = False
        while not rospy.is_shutdown():
            positions = []
            with self.lock:
                for idx, jid in enumerate(joint_ids):
                    try:
                        ang_deg = self.safe_dr_get_angle(jid)
                        pos = math.radians(ang_deg)
                    except Exception as e:
                        rospy.logwarn_throttle(2.0, f"读取关节 {jid} 失败: {e}")
                        pos = 0.0
                    positions.append(pos)

            # 平滑
            if not initialized:
                smooth_positions = positions[:6]
                initialized = True
            else:
                for i in range(6):
                    smooth_positions[i] = (1 - self.joint_states_smooth_alpha) * smooth_positions[i] + self.joint_states_smooth_alpha * positions[i]

            js = JointState()
            js.header.stamp = rospy.Time.now()
            js.name = joint_names
            js.position = smooth_positions
            self.joint_state_pub.publish(js)
            rate.sleep()

    # ---------- 获取当前关节值（供 IK seed 使用） ----------
    def get_current_joint_values(self):
        positions = []
        with self.lock:
            for jid in joint_ids:
                try:
                    ang_deg = self.safe_dr_get_angle(jid)
                    positions.append(math.radians(ang_deg))
                except:
                    positions.append(0.0)
        return positions

    # ---------- 安全读取 SDK ----------
    def safe_dr_get_angle(self, jid):
        return dr.get_angle(id_num=jid)

    # ---------- 回滚 ----------
    def rollback_pose(self, prev_pos, prev_q):
        self.target_pose.pose.position.x, \
        self.target_pose.pose.position.y, \
        self.target_pose.pose.position.z = prev_pos
        self.target_pose.pose.orientation = prev_q
        rospy.logwarn("⚠️ IK失败 → 回滚位姿")

    # ---------- 工具 ----------
    @staticmethod
    def quat_to_rot(q):
        return tft.quaternion_matrix([q.x, q.y, q.z, q.w])[:3, :3]

    @staticmethod
    def print_pose(pose, prefix="", throttle=None):
        p = pose.pose.position
        q = pose.pose.orientation
        yaw, pitch, roll = np.degrees(tft.euler_from_quaternion([q.x, q.y, q.z, q.w]))
        msg = f"{prefix}: x={p.x:.3f}, y={p.y:.3f}, z={p.z:.3f}, roll={roll:.1f}°, pitch={pitch:.1f}°, yaw={yaw:.1f}°"
        if throttle:
            rospy.loginfo_throttle(throttle, msg)
        else:
            rospy.loginfo(msg)


if __name__ == "__main__":
    try:
        TeleopPoseReal()
    except rospy.ROSInterruptException:
        pass
    finally:
        roscpp_shutdown()
